#include "linear_p2_crypto_agent.hpp"
#include "log_helper.hpp"
#include "math_util.hpp"
#include <map>

/*** 4*size*size <= 16384 => m = 16384 ***/
// * precision = 20; bits = 170
//     size   | noise   | err
//   - 4*4    | 138.138 | 0.00653967
//   - 8*8    | 138.284 | 0.0665066
//   - 16*16  | 141.356 | 0.605521
//   - 32*32  | 144.391 | 4.91075
//   - 64*64  | 145.437 | 8.94333
//
// * precision = 24; bits = 180
//     size   | noise   | err
//   - 64*64  | 160.436 | 0.650136
//
// * precision = 30; bits = 180
//     size   | noise   | err
//   - 4*4    | 177.113 | 7.34168e-06
//   - 8*8    | 180.271 | 7.16361e-05
//   - 16*16  | 180.43  | 0.000698855
//   - 32*32  | 183.465 | 0.00557298
//   - 64*64  | 184.51  | 0.0101348
//
/*** m = 4*128*128 = 65536 ***/
// *  = 20; bits = 558
//     size   | precision | noise   | err
//   128*128  | 20        | 380.533 | 38.1145
//   128*128  | 30        | 419.687 | 0.0101348
//
/*** m = 4*256*256 = 131072 ***/
// *  = 20; bits = 1098
//     size   | precision | noise   | err
//   256*256  | 20        | 1224.11 | 380.035
// clang-format off
std::map<long, std::vector<std::pair<long, long>>>
    ckks_table = {
        {16384,  {{119, 2}/* 157.866 */,  {170, 3}/* 119.080 */, {180, 4}/* 97.088 */, {230, 4}/* 88.7942 */, {407, 8}/* 46.4566: too low!! */}},
        {32768,  {{239, 2}/* 155.025 */,  {299, 3}/* 129.816 */,  {358, 6}/* 129.741 */}},
        {65536,  {{558, 2}/* 128.851 */,  {613, 3},  {669, 4},  {717, 6},  {725, 8}}},
        {131072, {{1098, 2}, {1255, 3}, {1329, 4}, {1387, 5}, {1435, 6}, {1445, 8}}},
        {262144, {{2234, 2}, {2511, 3}, {2646, 4}, {2763, 5}, {2870, 6}, {2940, 8}}}};
// clang-format on

static void find_bits_and_c(long &m, long &bits, long &c) {
  auto res = ckks_table.lower_bound(m);
  if (res == ckks_table.end()) {
    ERROR_PRINTF("m = %ld is too large for CKKS in HElib\n", m);
  }
  m = res->first;
  auto bits_c_table = res->second;
  for (auto &pair : bits_c_table) {
    if (bits <= pair.first) {
      bits = pair.first;
      c = pair.second;
      break;
    }
  }
  auto largest_bits = bits_c_table.back().first;
  ASSERT_PRINTF(
      bits <= bits_c_table.back().first,
      "bits=%ld > largest_bits=%ld: Not using recommended settings from "
      "`lib/helib/examples/tutorial/01_ckks_basics.cpp`\n",
      bits, largest_bits);
}

template <>
linearp2_hmmcs::linearp2_hmmcs(const params<helib::BGV> &params,
                               long expected_nslot) {
  // Hamming weight of the secret key
  const unsigned long hwt = params.hwt;
  // Number of columns of Key-Switching matrix (default = 2 or 3)
  const unsigned long ksm_ncol = params.ksm_ncol;
  // Number of modulus_chain_bits of the modulus chain
  const unsigned long modulus_chain_bits = params.modulus_chain_bits;
  // Plaintext prime modulus
  long t;
  // Bit length low bound of plaintext prime modulus
  const unsigned long t_bits = params.t_bits;
  // Hensel lifting (default = 1)
  const unsigned long r = 1;

  setup_time_ -= omp_get_wtime();       //! Setup - Start
  find_params_time_ -= omp_get_wtime(); //! Find m & t - Start
  long next_power_of2 =
      std::pow(2, static_cast<int>(std::ceil(std::log2(expected_nslot))));
  //! For power-of-2, we find 2 \times (\phi(m)/2) "hypercube"
  //! To perform matrix multiplication, we must let (\phi(m)/2) = |maxtrix|.
  //! It means that we waste a column of slots
  //! FIXME:
  // long m = 2l * next_power_of2;
  long m = 4l * next_power_of2;
  while (m < params.m_lower_bound) {
    next_power_of2 *= 2l;
    m *= 2l;
  }
  std::vector<long> candidates = {2l, next_power_of2};
  find_ti_auto(t, t_bits, m, candidates);

  std::vector<long> gens;
  std::vector<long> ords;
  find_generators(gens, ords, m, t, candidates);

  find_params_time_ -= omp_get_wtime(); //! Find m & t - Stop

  build_context_time_ -= omp_get_wtime(); //! build context - Start
  context_ = std::shared_ptr<helib::Context>(helib::ContextBuilder<helib::BGV>()
                                                 .m(m)
                                                 .p(t)
                                                 .r(r)
                                                 .bits(modulus_chain_bits)
                                                 .skHwt(hwt)
                                                 .gens(gens)
                                                 .ords(ords)
                                                 .c(ksm_ncol)
                                                 .buildPtr());
  build_context_time_ += omp_get_wtime(); //! build context - Stop
  setup_time_ += omp_get_wtime();         //! Setup - Stop
}

template <>
linearp2_hmmcs::linearp2_hmmcs(const params<helib::CKKS> &params,
                               long expected_nslot) {
  // the number of bits of precision
  // when data is encoded, encrypted, or decrypted
  const unsigned long precision = params.precision;
  ASSERT_PRINTF(
      precision <= 40,
      "It is not recommended to set precision greater than 40 or so\n");
  // Number of bits of the ciphertext modulus
  long bits = params.bits;
  // Number of columns of Key-Switching matrix (default = 2 or 3)
  long ksm_ncol;

  setup_time_ -= omp_get_wtime();       //! Setup - Start
  find_params_time_ -= omp_get_wtime(); //! Find m & bits - Start
  long next_power_of2 =
      std::pow(2, static_cast<int>(std::ceil(std::log2(expected_nslot))));
  long m = 4l * next_power_of2;
  while (m < params.m_lower_bound) {
    m *= 2l;
  }
  find_bits_and_c(m, bits, ksm_ncol);
  find_params_time_ += omp_get_wtime(); //! Find m & bits - Stop

  build_context_time_ -= omp_get_wtime(); //! build context - Start
  context_ =
      std::shared_ptr<helib::Context>(helib::ContextBuilder<helib::CKKS>()
                                          .m(m)
                                          .precision(precision)
                                          .bits(bits)
                                          .c(ksm_ncol)
                                          .buildPtr());
  build_context_time_ += omp_get_wtime(); //! build context - Stop
  setup_time_ += omp_get_wtime();         //! Setup - Stop
}

void linearp2_hmmcs::report_time() const {
  // clang-format off
  std::cout << "* Setup Time                 : " << setup_time_ << " s\n";
  std::cout << "  * Find params              : " << find_params_time_ * 1000.0 << " ms\n";
  std::cout << "  * Build context            : " << build_context_time_ << " s\n";
  // clang-format on
}

template <typename Scheme>
linearp2_hmmcc<Scheme>::linearp2_hmmcc(const params<Scheme> &params,
                                       long expected_nslot)
    : linearp2_hmmcs(params, expected_nslot),
      hmm_crypto_client<Scheme>(*context_) {
  //* Construct `linearp2_hmmcc` instance to build the HElib context
  //* Using `context` to construct `hmm_crypto_client` instance

  this->setup_time_ += this->key_gen_time_; //! Part of Setup
}

template <typename Scheme>
linearp2_hmmcc<Scheme>::linearp2_hmmcc(const linearp2_hmmcc<Scheme> &hmmcs)
    : linearp2_hmmcs(hmmcs.context(), hmmcs.build_context_time()),
      hmm_crypto_client<Scheme>(*context_) {}

template <typename Scheme> void linearp2_hmmcc<Scheme>::report_time() const {
  linearp2_hmmcs::report_time();
  std::cout << "  * KeyGen Time              : " << this->key_gen_time_
            << " s\n";
}

//* Encryption & Decryption
static inline void pack_cube(helib::Ptxt<helib::BGV> &ptxt,
                             const NTL::mat_ZZ &mat,
                             const helib::EncryptedArray &ea) {
  const long row = mat.NumRows();
  const long col = mat.NumCols();

  bool row_is_2 = ea.dimension() == 2 && ea.sizeOfDimension(0) == 2;
  ASSERT_PRINTF(row_is_2 || ea.dimension() == 1,
                "Invalid linear packing context\n");

  for (auto i = 0; i < row; i++) {
    for (auto j = 0; j < col; j++) {
      int idx = i * col + j;
      ptxt[idx] = NTL::conv<long>(mat[i][j]);
    }
  }
}

template <>
void linearp2_hmmcc<helib::BGV>::encrypt(helib::Ctxt &ctxt,
                                         const NTL::mat_ZZ &mat, ENC type) {
  // sanity check
  ASSERT_PRINTF(ctxt.getPubKey() == this->public_key(),
                "Encrypt: public key mismatch\n");

  // Create a vector of long with nslots elements
  helib::Ptxt<helib::BGV> ptxt(*context_);

  // Packing matrix
  pack_cube(ptxt, mat, context_->getEA());

  hmm_crypto_client::encrypt(ctxt, ptxt, type);
}

template <>
void linearp2_hmmcc<helib::CKKS>::encrypt(helib::Ctxt &ctxt,
                                          const NTL::mat_ZZ &mat, ENC type) {
  // sanity check
  ASSERT_PRINTF(ctxt.getPubKey() == this->public_key(),
                "Encrypt: public key mismatch\n");

  // Create a vector of long with nslots elements
  std::vector<double> slots(context_->getEA().size(), 0.0);

  // Packing matrix
  const long row = mat.NumRows();
  const long col = mat.NumCols();
  for (auto i = 0; i < row; i++) {
    for (auto j = 0; j < col; j++) {
      int idx = i * col + j;
      slots[idx] = NTL::conv<long>(mat[i][j]);
    }
  }

  helib::PtxtArray ptxt(*context_);
  ptxt.load(slots);
  ptxt.encrypt(ctxt);
}

template <typename Scheme>
helib::Ctxt linearp2_hmmcc<Scheme>::encrypt(const NTL::mat_ZZ &mat, ENC type) {
  auto ctxt = this->get_empty_ctxt();
  this->encrypt(ctxt, mat, type);
  return std::move(ctxt);
}

template <typename Scheme>
helib::Ctxt linearp2_hmmcc<Scheme>::encrypt(const std::vector<double> &vec,
                                            ENC type) {
  auto ctxt = this->get_empty_ctxt();
  helib::PtxtArray ptxt(*context_, vec);
  ptxt.encrypt(ctxt);
  return std::move(ctxt);
}

template <typename Scheme>
helib::Ctxt linearp2_hmmcc<Scheme>::encrypt(const std::vector<long> &vec,
                                            ENC type) {
  auto ctxt = this->get_empty_ctxt();
  helib::PtxtArray ptxt(*context_, vec);
  ptxt.encrypt(ctxt);
  return std::move(ctxt);
}

template <>
void linearp2_hmmcc<helib::BGV>::decrypt(NTL::mat_ZZ &mat,
                                         const helib::Ctxt &cmat) {
  // Create a vector of long with nslots elements
  helib::Ptxt<helib::BGV> ptxt(*context_);
  // Decrypt
  hmm_crypto_client::decrypt(ptxt, cmat);

  const helib::EncryptedArray &ea = context_->getEA();

  bool row_is_2 = ea.dimension() == 2 && ea.sizeOfDimension(0) == 2;
  long dim_idx = row_is_2 ? 1 : 0;

  long nslot = ea.sizeOfDimension(dim_idx);
  mat.SetDims(1, nslot);
  for (auto i = 0; i < nslot; i++) {
    mat[0][i] = NTL::coeff(ptxt[i].getData(), 0);
  }
}

template <typename Scheme>
void linearp2_hmmcc<Scheme>::decrypt(std::vector<double> &vec,
                                     const helib::Ctxt &cmat) {
  // Create a vector of long with nslots elements
  helib::PtxtArray ptxt(*context_);
  // Decrypt
  ptxt.decrypt(cmat, *(this->sk_));
  ptxt.store(vec);
}

template <typename Scheme>
void linearp2_hmmcc<Scheme>::decrypt(std::vector<long> &vec,
                                     const helib::Ctxt &cmat) {
  if (std::is_same<Scheme, helib::CKKS>::value) {
    INFO_PRINTF("There is a truncation error\n");
  }
  // Create a vector of long with nslots elements
  helib::PtxtArray ptxt(*context_);
  // Decrypt
  ptxt.decrypt(cmat, *(this->sk_));
  ptxt.store(vec);
}

template <>
void linearp2_hmmcc<helib::CKKS>::decrypt(NTL::mat_ZZ &mat,
                                          const helib::Ctxt &cmat) {
  INFO_PRINTF("There is a truncation error\n");
  // Create a vector of long with nslots elements
  helib::Ptxt<helib::CKKS> ptxt(*context_);
  // Decrypt
  hmm_crypto_client::decrypt(ptxt, cmat);

  const helib::EncryptedArray &ea = context_->getEA();

  bool row_is_2 = ea.dimension() == 2 && ea.sizeOfDimension(0) == 2;
  long dim_idx = row_is_2 ? 1 : 0;

  long nslot = ea.sizeOfDimension(dim_idx);
  mat.SetDims(1, nslot);
  for (auto i = 0; i < nslot; i++) {
    mat[0][i] = ptxt[i].real();
  }
}

// //* template specialization
// for BGV
template class linearp2_hmmcc<helib::BGV>;
// for CKKS
template class linearp2_hmmcc<helib::CKKS>;